package com.huawei.admins.platform.gateway.core.util;

import com.huawei.admins.platform.gateway.support.HttpStatusHolder;
import io.netty.buffer.Unpooled;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.core.io.buffer.DefaultDataBuffer;
import org.springframework.core.io.buffer.NettyDataBuffer;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.AbstractServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.util.Assert;
import org.springframework.web.reactive.DispatcherHandler;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.util.UriComponentsBuilder;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.net.URI;
import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedHashSet;
import java.util.Map;
import java.util.Set;
import java.util.function.Function;
import java.util.function.Predicate;

public class ServerWebExchangeUtils {

    private static final Logger logger = LoggerFactory.getLogger(ServerWebExchangeUtils.class);

    /**
     * Preserve-Host header attribute name.
     */
    public static final String PRESERVE_HOST_HEADER_ATTRIBUTE = qualify("preserveHostHeader");

    /**
     * URI template variables attribute name.
     */
    public static final String URI_TEMPLATE_VARIABLES_ATTRIBUTE = qualify("uriTemplateVariables");

    /**
     * Client response attribute name.
     */
    public static final String CLIENT_RESPONSE_ATTR = qualify("gatewayClientResponse");

    /**
     * Client response connection attribute name.
     */
    public static final String CLIENT_RESPONSE_CONN_ATTR = qualify("gatewayClientResponseConnection");

    /**
     * Client response header names attribute name.
     */
    public static final String CLIENT_RESPONSE_HEADER_NAMES = qualify("gatewayClientResponseHeaderNames");

    /**
     * Gateway route attribute name.
     */
    public static final String GATEWAY_ROUTE_ATTR = qualify("gatewayRoute");

    /**
     * Original Reactor Context corresponding to the processed request.
     */
    public static final String GATEWAY_REACTOR_CONTEXT_ATTR = qualify("gatewayReactorContext");

    /**
     * Gateway request URL attribute name.
     */
    public static final String GATEWAY_REQUEST_URL_ATTR = qualify("gatewayRequestUrl");

    /**
     * Gateway original request URL attribute name.
     */
    public static final String GATEWAY_ORIGINAL_REQUEST_URL_ATTR = qualify("gatewayOriginalRequestUrl");

    /**
     * Gateway handler mapper attribute name.
     */
    public static final String GATEWAY_HANDLER_MAPPER_ATTR = qualify("gatewayHandlerMapper");

    /**
     * Gateway scheme prefix attribute name.
     */
    public static final String GATEWAY_SCHEME_PREFIX_ATTR = qualify("gatewaySchemePrefix");

    /**
     * Gateway predicate route attribute name.
     */
    public static final String GATEWAY_PREDICATE_ROUTE_ATTR = qualify("gatewayPredicateRouteAttr");

    /**
     * Gateway predicate matched path attribute name.
     */
    public static final String GATEWAY_PREDICATE_MATCHED_PATH_ATTR = qualify("gatewayPredicateMatchedPathAttr");

    /**
     * Gateway predicate matched path route id attribute name.
     */
    public static final String GATEWAY_PREDICATE_MATCHED_PATH_ROUTE_ID_ATTR = qualify(
            "gatewayPredicateMatchedPathRouteIdAttr");

    /**
     * Gateway predicate path container attribute name.
     */
    public static final String GATEWAY_PREDICATE_PATH_CONTAINER_ATTR = qualify("gatewayPredicatePathContainer");

    /**
     * Weight attribute name.
     */
    public static final String WEIGHT_ATTR = qualify("routeWeight");

    /**
     * Original response Content-Type attribute name.
     */
    public static final String ORIGINAL_RESPONSE_CONTENT_TYPE_ATTR = "original_response_content_type";

    /**
     * CircuitBreaker execution exception attribute name.
     */
    public static final String CIRCUITBREAKER_EXECUTION_EXCEPTION_ATTR = qualify("circuitBreakerExecutionException");

    /**
     * Used when a routing filter has been successfully called. Allows users to write
     * custom routing filters that disable built in routing filters.
     */
    public static final String GATEWAY_ALREADY_ROUTED_ATTR = qualify("gatewayAlreadyRouted");

    /**
     * Gateway already prefixed attribute name.
     */
    public static final String GATEWAY_ALREADY_PREFIXED_ATTR = qualify("gatewayAlreadyPrefixed");

    /**
     * Cached ServerHttpRequestDecorator attribute name. Used when
     */
    public static final String CACHED_SERVER_HTTP_REQUEST_DECORATOR_ATTR = "cachedServerHttpRequestDecorator";

    /**
     * Cached request body key. Used when
     */
    public static final String CACHED_REQUEST_BODY_ATTR = "cachedRequestBody";

    public static final String GATEWAY_LOADBALANCER_RESPONSE_ATTR = qualify("gatewayLoadBalancerResponse");

    /**
     * Gateway Client {@code Observation} attribute name.
     */
    public static final String GATEWAY_OBSERVATION_ATTR = qualify("gateway.observation");

    private static final byte[] EMPTY_BYTES = {};

    private ServerWebExchangeUtils() {
        throw new AssertionError("Must not instantiate utility class.");
    }

    private static String qualify(String attr) {
        return ServerWebExchangeUtils.class.getName() + "." + attr;
    }

    public static void setAlreadyRouted(ServerWebExchange exchange) {
        exchange.getAttributes().put(GATEWAY_ALREADY_ROUTED_ATTR, true);
    }

    public static void removeAlreadyRouted(ServerWebExchange exchange) {
        exchange.getAttributes().remove(GATEWAY_ALREADY_ROUTED_ATTR);
    }

    public static HttpStatus parse(String statusString) {
        HttpStatus httpStatus;

        try {
            int status = Integer.parseInt(statusString);
            httpStatus = HttpStatus.resolve(status);
        }
        catch (NumberFormatException e) {
            // try the enum string
            httpStatus = HttpStatus.valueOf(statusString.toUpperCase());
        }
        return httpStatus;
    }

    public static void addOriginalRequestUrl(ServerWebExchange exchange, URI url) {
        exchange.getAttributes().computeIfAbsent(GATEWAY_ORIGINAL_REQUEST_URL_ATTR, s -> new LinkedHashSet<>());
        LinkedHashSet<URI> uris = exchange.getRequiredAttribute(GATEWAY_ORIGINAL_REQUEST_URL_ATTR);
        uris.add(url);
    }

    public static boolean setResponseStatus(ServerWebExchange exchange, HttpStatusHolder statusHolder) {
        if (exchange.getResponse().isCommitted()) {
            return false;
        }
        if (logger.isDebugEnabled()) {
            logger.debug("Setting response status to {}", statusHolder);
        }
        if (statusHolder.getHttpStatus() != null) {
            return setResponseStatus(exchange, statusHolder.getHttpStatus());
        }
        if (statusHolder.getStatus() != null && exchange.getResponse() instanceof AbstractServerHttpResponse) { // non-standard
            exchange.getResponse().setRawStatusCode(statusHolder.getStatus());
            return true;
        }
        return false;
    }

    public static boolean setResponseStatus(ServerWebExchange exchange, HttpStatus httpStatus) {
        boolean response = exchange.getResponse().setStatusCode(httpStatus);
        if (!response && logger.isWarnEnabled()) {
            logger.warn("Unable to set status code to {}. Response already committed.", httpStatus);
        }
        return response;
    }

    public static void reset(ServerWebExchange exchange) {
        Set<String> addedHeaders = exchange.getAttributeOrDefault(CLIENT_RESPONSE_HEADER_NAMES, Collections.emptySet());
        addedHeaders.forEach(header -> exchange.getResponse().getHeaders().remove(header));
        removeAlreadyRouted(exchange);
    }
}
